#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2024/3/11 14:10
# @Author  : jeff
# @FileName: user_service
import datetime
import json
import uuid

from flask import request
from loguru import logger

from src.config.redis_cache import RedisCache
from src.enums.status_enum import Status
from src.model.scan_task import ScanTask
from src.model.scan_work import ScanWork
from src.service.base_service import BaseService


class AssetService(BaseService):
    @staticmethod
    def task_add():
        params_str = request.json.get("params")
        execute_time_str = request.json.get("execute_time")
        # 将params字符串解析为字典
        if isinstance(params_str, str):
            params = json.loads(params_str)
        else:
            params = params_str

        if isinstance(execute_time_str, str):
            execute_time = json.loads(execute_time_str)
        else:
            execute_time = execute_time_str

        callback_url = request.json.get("callback_url")  # 修改变量名为小写加下划线
        progress_url = request.json.get("progress_url", "http")

        logger.info(f'params:{params},callback_url:{callback_url}')  # 修改日志输出中的变量名

        # 创建scanwork对象并对其中的值赋值
        data = {
            "params": params,
            "callback_url": callback_url,
            "callback_type": 'API',
            "callback_status": 'pending',
            "execute_time": execute_time,
            "uuid": str(uuid.uuid4()),  # 假设使用uuid库生成唯一标识符
            "status": 'pending',  # 假设初始状态为'PENDING'
            "source": 'default',  # 根据实际情况设置默认值或从其他地方获取
            "work_type": 'portscan',  # 根据实际情况设置默认值或从其他地方获取
            "retry": 0,  # 假设初始重试次数为0
            "priority": 1,  # 假设默认优先级为1
            "progress": 0,  # 假设初始进度为0
            "progress_type": 'API',  # 根据实际情况设置默认值或从其他地方获取
            "progress_url": progress_url,  # 根据实际情况设置默认值或从其他地方获取
            "create_time": datetime.datetime.now(),  # 使用当前时间作为创建时间
            "update_time": datetime.datetime.now()  # 使用当前时间作为更新时间
        }
        new_data = ScanWork(**data)
        # 创建scanwork对象并对其中的值赋值
        scan_work = ScanWork.insert(new_data)
        if not scan_work:
            return Status.ERROR_SELECT, {"message": "Internal Server Error"}
        # {"ip": "192.168.31.0,192.168.31.1,192.168.31.2,192.168.31.3,192.168.31.4,192.168.31.5,192.168.31.6,192.168.31.7,192.168.31.8,192.168.31.9,192.168.31.10,192.168.31.11,192.168.31.12,192.168.31.13,192.168.31.14,192.168.31.15,192.168.31.16,192.168.31.17,192.168.31.18,192.168.31.19,192.168.31.20,192.168.31.21,192.168.31.22,192.168.31.23,192.168.31.24,192.168.31.25,192.168.31.26,192.168.31.27,192.168.31.28,192.168.31.29,192.168.31.30,192.168.31.31,192.168.31.32,192.168.31.33,192.168.31.34,192.168.31.35,192.168.31.36,192.168.31.37,192.168.31.38,192.168.31.39,192.168.31.40,192.168.31.41,192.168.31.42,192.168.31.43,192.168.31.44,192.168.31.45,192.168.31.46,192.168.31.47,192.168.31.48,192.168.31.49,192.168.31.50,192.168.31.51,192.168.31.52,192.168.31.53,192.168.31.54,192.168.31.55,192.168.31.56,192.168.31.57,192.168.31.58,192.168.31.59,192.168.31.60,192.168.31.61,192.168.31.62,192.168.31.63,192.168.31.64,192.168.31.65,192.168.31.66,192.168.31.67,192.168.31.68,192.168.31.69,192.168.31.70,192.168.31.71,192.168.31.72,192.168.31.73,192.168.31.74,192.168.31.75,192.168.31.76,192.168.31.77,192.168.31.78,192.168.31.79,192.168.31.80,192.168.31.81,192.168.31.82,192.168.31.83,192.168.31.84,192.168.31.85,192.168.31.86,192.168.31.87,192.168.31.88,192.168.31.89,192.168.31.90,192.168.31.91,192.168.31.92,192.168.31.93,192.168.31.94,192.168.31.95,192.168.31.96,192.168.31.97,192.168.31.98,192.168.31.99,192.168.31.100,192.168.31.101,192.168.31.102,192.168.31.103,192.168.31.104,192.168.31.105,192.168.31.106,192.168.31.107,192.168.31.108,192.168.31.109,192.168.31.110,192.168.31.111,192.168.31.112,192.168.31.113,192.168.31.114,192.168.31.115,192.168.31.116,192.168.31.117,192.168.31.118,192.168.31.119,192.168.31.120,192.168.31.121,192.168.31.122,192.168.31.123,192.168.31.124,192.168.31.125,192.168.31.126,192.168.31.127,192.168.31.128,192.168.31.129,192.168.31.130,192.168.31.131,192.168.31.132,192.168.31.133,192.168.31.134,192.168.31.135,192.168.31.136,192.168.31.137,192.168.31.138,192.168.31.139,192.168.31.140,192.168.31.141,192.168.31.142,192.168.31.143,192.168.31.144,192.168.31.145,192.168.31.146,192.168.31.147,192.168.31.148,192.168.31.149,192.168.31.150,192.168.31.151,192.168.31.152,192.168.31.153,192.168.31.154,192.168.31.155,192.168.31.156,192.168.31.157,192.168.31.158,192.168.31.159,192.168.31.160,192.168.31.161,192.168.31.162,192.168.31.163,192.168.31.164,192.168.31.165,192.168.31.166,192.168.31.167,192.168.31.168,192.168.31.169,192.168.31.170,192.168.31.171,192.168.31.172,192.168.31.173,192.168.31.174,192.168.31.175,192.168.31.176,192.168.31.177,192.168.31.178,192.168.31.179,192.168.31.180,192.168.31.181,192.168.31.182,192.168.31.183,192.168.31.184,192.168.31.185,192.168.31.186,192.168.31.187,192.168.31.188,192.168.31.189,192.168.31.190,192.168.31.191,192.168.31.192,192.168.31.193,192.168.31.194,192.168.31.195,192.168.31.196,192.168.31.197,192.168.31.198,192.168.31.199,192.168.31.200,192.168.31.201,192.168.31.202,192.168.31.203,192.168.31.204,192.168.31.205,192.168.31.206,192.168.31.207,192.168.31.208,192.168.31.209,192.168.31.210,192.168.31.211,192.168.31.212,192.168.31.213,192.168.31.214,192.168.31.215,192.168.31.216,192.168.31.217,192.168.31.218,192.168.31.219,192.168.31.220,192.168.31.221,192.168.31.222,192.168.31.223,192.168.31.224,192.168.31.225,192.168.31.226,192.168.31.227,192.168.31.228,192.168.31.229,192.168.31.230,192.168.31.231,192.168.31.232,192.168.31.233,192.168.31.234,192.168.31.235,192.168.31.236,192.168.31.237,192.168.31.238,192.168.31.239,192.168.31.240,192.168.31.241,192.168.31.242,192.168.31.243,192.168.31.244,192.168.31.245,192.168.31.246,192.168.31.247,192.168.31.248,192.168.31.249,192.168.31.250,192.168.31.251,192.168.31.252,192.168.31.253,192.168.31.254,192.168.31.255", "port": "top1000", "scan_type": "CONNECT"}
        # 拆分扫描任务，如果包含的IP太多的话进行任务拆分
        # 拆分扫描任务，如果包含的IP太多的话进行任务拆分
        ips = params.get('ip', '').split(',')
        port_str = params.get('port', '')
        num = 256
        if len(ips) > num:
            for i in range(0, len(ips), num):
                task_data = {
                    "execute_time": execute_time,
                    "work_uuid": scan_work.uuid,  # 假设使用uuid库生成唯一标识符
                    "uuid": str(uuid.uuid4()),
                    "status": 'pending',  # 假设初始状态为'PENDING'
                    "retry": 1,
                    "priority": 1,
                    "create_time": datetime.datetime.now(),  # 使用当前时间作为创建时间
                    "update_time": datetime.datetime.now()  # 使用当前时间作为更新时间
                }
                task_data['params'] = {
                    "ip": ips[i:i + num],
                    "port": port_str
                }
                new_ips_data = ScanTask(**task_data)
                ScanTask.insert(new_ips_data)
        else:
            task_data = {
                "execute_time": execute_time,
                "work_uuid": str(scan_work.uuid),  # 假设使用uuid库生成唯一标识符
                "uuid": str(uuid.uuid4()),
                "status": 'pending',  # 假设初始状态为'PENDING'
                "retry": 1,
                "priority": 1,
                "create_time": datetime.datetime.now(),  # 使用当前时间作为创建时间
                "update_time": datetime.datetime.now()  # 使用当前时间作为更新时间
            }
            task_data['params'] = {
                "ip": ips,
                "port": port_str
            }
            new_ips_data = ScanTask(**task_data)
            ScanTask.insert(new_ips_data)
        return Status.OK, {"data": scan_work.id}

    #下发网站爬取相关的任务时会用到此方法
    @staticmethod
    def get_web_info():
        urls = request.json.get("urls")
        if urls:
            redis_cache = RedisCache()
            redis_cache.lpush('website:info:pending', urls);
